#include "histogram_equalization.h"

#include "../histogram.h"
#include <memory>

HistogramEqualization::HistogramEqualization(PNM* img) :
	Transformation(img)
{
}

HistogramEqualization::HistogramEqualization(PNM* img, ImageViewer* iv) :
	Transformation(img, iv)
{
}

PNM* HistogramEqualization::transform()
{
	int width = image->width();
	int height = image->height();

	PNM* newImage = new PNM(width, height, image->format());

	if (image->format() == QImage::Format_Indexed8)
	{
		auto l_hist = image->getHistogram()->get(Histogram::LChannel);

		unique_ptr<QHash<int, float>> D(new QHash<int, float>);
		for (int i = 0; i < 256; i++)
		{
			float p = (float)(l_hist->value(i)) / (float)(width*height);
			D->insert(i, p);
		}

		for (int x=0; x<width; x++)
			for (int y=0; y<height; y++)
			{
				QRgb pixel = image->pixel(x,y);

				int v = qGray(pixel);

				float sum = 0;
				for (int i = 0; i < v; i++)
				{
					sum += D->value(i);
				}

				int newColor((int)(sum*255.0f));				
				newImage->setPixel(x, y, newColor);
			}
	}
	else //if (image->format() == QImage::Format_RGB32)
	{
		auto r_hist = image->getHistogram()->get(Histogram::RChannel);
		auto g_hist = image->getHistogram()->get(Histogram::GChannel);
		auto b_hist = image->getHistogram()->get(Histogram::BChannel);


		unique_ptr<QHash<int, float>> DR(new QHash<int, float>);
		unique_ptr<QHash<int, float>> DG(new QHash<int, float>);
		unique_ptr<QHash<int, float>> DB(new QHash<int, float>);

		for (int i = 0; i < 256; i++)
		{
			float p = (float)(r_hist->value(i)) / (float)(width*height);
			DR->insert(i, p);

			p = (float)(g_hist->value(i)) / (float)(width*height);
			DG->insert(i, p);

			p = (float)(b_hist->value(i)) / (float)(width*height);
			DB->insert(i, p);
		}

		for (int x=0; x<width; x++)
			for (int y=0; y<height; y++)
			{
				QRgb pixel = image->pixel(x,y);

				int r = qRed(pixel);
				int g = qGreen(pixel);
				int b = qBlue(pixel);

				float sumr = 0;
				float sumg = 0;
				float sumb = 0;

				for (int i = 0; i < r; i++) sumr += DR->value(i);
				for (int i = 0; i < g; i++) sumg += DG->value(i);
				for (int i = 0; i < b; i++) sumb += DB->value(i);

				int newr = (int)(sumr*255.0f);
				int newg = (int)(sumg*255.0f);
				int newb = (int)(sumb*255.0f);

				QColor newColor(
					newr,
					newg,
					newb
				);

				newImage->setPixel(x,y,newColor.rgb());
			}
	}

	return newImage;
}

